﻿#include <Windows.h>
#include <iostream>
#include <vector>

#include "../ProcessProtectCommon/ProcessProtectCommo.h"

int
Error(const char* msg);

int
PrintUsage();

std::vector<DWORD> 
ParsePids(const wchar_t* buffer[], int count);

int 
wmain(int argc, const wchar_t* argv[])
{
    if (argc < 2) {
		return PrintUsage();
    }

	enum class Options {
		Unknown,
		Add, 
		Remove, 
		Clear,
		Query
	};

	Options Option;
	if (::_wcsicmp(argv[1], L"add") == 0) {
		Option = Options::Add;
	} else if (::_wcsicmp(argv[1], L"remove") == 0) {
		Option = Options::Remove;
	} else if (::_wcsicmp(argv[1], L"clear") == 0) {
		Option = Options::Clear;
	} else if (::_wcsicmp(argv[1], L"Query") == 0) {
		Option = Options::Query;
	} else {
		printf("Unknown option.\n");
		return PrintUsage();
	}

	HANDLE hFile = ::CreateFile(L"\\\\.\\" PROCESS_PROTECT_NAME, GENERIC_WRITE | GENERIC_READ, 0, nullptr, OPEN_EXISTING, 0, nullptr);
	if (hFile == INVALID_HANDLE_VALUE) {
		return Error("Failed to open device");
	}

	std::vector<DWORD> Pids;
	BOOL bSuccess = FALSE;
	DWORD dwBytes = 0;
	ULONG ulQueryPids[g_ulMaxPids];
	switch (Option) {
	case Options::Add: {
		Pids = ParsePids(argv + 2, argc - 2);
		bSuccess = ::DeviceIoControl(hFile, IOCTL_PROCESS_PROTECT_BY_PID, Pids.data(), static_cast<DWORD>(Pids.size() * sizeof(DWORD)), nullptr, 0, &dwBytes, nullptr);
		break;
	}

	case Options::Remove: {
		Pids = ParsePids(argv + 2, argc - 2);
		bSuccess = ::DeviceIoControl(hFile, IOCTL_PROCESS_UNPROTECT_BY_PID, Pids.data(), static_cast<DWORD>(Pids.size() * sizeof(DWORD)), nullptr, 0, &dwBytes, nullptr);
		break;
	}

	case Options::Clear: {
		bSuccess = ::DeviceIoControl(hFile, IOCTL_PROCESS_PROTECT_CLEAR, nullptr, 0, nullptr, 0, &dwBytes, nullptr);
		break;
	}

	case Options::Query: {
		bSuccess = ::DeviceIoControl(hFile, IOCTL_PROCESS_QUERY_PIDS, nullptr, 0, ulQueryPids, g_ulMaxPids * sizeof(ULONG), &dwBytes, nullptr);
		break;
	}
	}

	if (bSuccess == FALSE) {
		return Error("Failed in DeviceIoControl");
	}

	if (Option == Options::Query) {
		printf("Pids: ");
		for (int i = 0; i < g_ulMaxPids; i++) {
			if (ulQueryPids[i] != 0) {
				printf("%u ", ulQueryPids[i]);
			}
		}
		printf("\n");
	}

	printf("Operation succeeded.\n");

	::CloseHandle(hFile);

    return 0;
}

int 
Error(const char* msg)
{
	printf("%s (Error: %d)\n", msg, ::GetLastError());
	return 1;
}

int 
PrintUsage()
{
	printf("Protect [add | remove | clear] [pid] ...\n");
	return 0;
}

std::vector<DWORD>
ParsePids(const wchar_t* buffer[], int count)
{
	std::vector<DWORD> pids;
	for (int i = 0; i < count; i++) {
		pids.push_back(::_wtoi(buffer[i]));
	}
	return pids;
}